/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.alibaba.cloud.ai.example.chat.dashscope.controller;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;

/**
 * @author yuluo
 * @author <a href="mailto:yuluo08290126@gmail.com">yuluo</a>
 */

@RestController
@RequestMapping("/model")
public class DashScopeChatModelController {

    private static final String DEFAULT_PROMPT = "你好，介绍下你自己吧。";

    private final ChatModel dashScopeChatModel;

    public DashScopeChatModelController(ChatModel chatModel) {
        this.dashScopeChatModel = chatModel;
    }

    /**
     * 最简单的使用方式，没有任何 LLMs 参数注入。
     *
     * @return String types.
     */
    @GetMapping("/simple/chat")
    public String simpleChat() {
        return dashScopeChatModel.call(
                        new Prompt(
                                DEFAULT_PROMPT,
                                DashScopeChatOptions.builder()
                                        .withModel(DashScopeApi.ChatModel.QWEN_PLUS.getValue())
                                        .build())
                ).getResult()
                .getOutput()
                .getText();
    }

    /**
     * Stream 流式调用。可以使大模型的输出信息实现打字机效果。
     *
     * @return Flux<String> types.
     */
    @GetMapping("/stream/chat")
    public Flux<String> streamChat(HttpServletResponse response) {

        // 避免返回乱码
        response.setCharacterEncoding("UTF-8");

        Flux<ChatResponse> stream = dashScopeChatModel.stream(
                new Prompt(
                        DEFAULT_PROMPT,
                        DashScopeChatOptions.builder()
                                .withModel(DashScopeApi.ChatModel.QWEN_PLUS.getValue())
                                .build()
                )
        );
        return stream.map(resp -> resp.getResult().getOutput().getText());
    }

    /**
     * 演示如何获取 LLM 得 token 信息
     */
    @GetMapping("/tokens")
    public Map<String, Object> tokens(HttpServletResponse response) {
        ChatResponse chatResponse = dashScopeChatModel.call(
                new Prompt(
                        DEFAULT_PROMPT, DashScopeChatOptions.builder()
                        .withModel(DashScopeApi.ChatModel.QWEN_PLUS.getValue())
                        .build()
                )
        );

        Map<String, Object> res = new HashMap<>();
        res.put("output", chatResponse.getResult().getOutput().getText());
        res.put("output_token", chatResponse.getMetadata().getUsage().getCompletionTokens());
        res.put("input_token", chatResponse.getMetadata().getUsage().getPromptTokens());
        res.put("total_token", chatResponse.getMetadata().getUsage().getTotalTokens());

        return res;
    }
    /*
        {
            "output": "你好！我是Qwen3，阿里巴巴集团旗下的通义千问大模型。我是一个超大规模语言模型，能够回答问题、创作文字，比如写故事、写公文、写邮件、写剧本、逻辑推理、编程等等，还能表达观点，玩游戏等。我支持多种语言，包括但不限于中文、英文、德语、法语、西班牙语等。如果你有任何问题或需要帮助，欢迎随时告诉我！",
            "output_token": 90,
            "input_token": 19,
            "total_token": 109
        }
    */


    /**
     * 使用编程方式自定义 LLMs ChatOptions 参数， {@link com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions}
     * 优先级高于在 application.yml 中配置的 LLMs 参数！
     */
    @GetMapping("/custom/chat")
    public String customChat() {
        DashScopeChatOptions customOptions = DashScopeChatOptions.builder()
                .withTopP(0.7)
                .withTopK(50)
                .withTemperature(0.8)
                .build();

        return dashScopeChatModel.call(new Prompt(DEFAULT_PROMPT, customOptions)).getResult().getOutput().getText();
    }


    // 如果体验 web search 和 自定义请求头，请本地编译主干仓库。

    /**
     * DashScope 联网搜索功能演示
     * 参数：https://help.aliyun.com/zh/model-studio/use-qwen-by-calling-api
     */
    @GetMapping("/dashscope/web-search")
    public Flux<String> dashScopeWebSearch(HttpServletResponse response) {

        // 定义要发送给模型的提示信息，请求其搜索关于 Spring AI 的相关介绍
        String prompt = "搜索下关于 Spring AI 的介绍";
        // 设置响应字符编码为 UTF-8，防止中文乱码
        response.setCharacterEncoding("UTF-8");

        // 构建 DashScope 搜索选项，配置搜索行为相关参数
        var searchOptions = DashScopeApi.SearchOptions.builder()
                .forcedSearch(true) // 强制启用联网搜索
                .enableSource(true) // 启用返回搜索结果来源链接
                .searchStrategy("pro") // 使用高级搜索策略
                .enableCitation(true) // 启用引用标记
                .citationFormat("[<number>]") // 设置引用格式为带方括号的数字编号
                .build();

        // 构建 ChatModel 所需的选项，包含模型名称、搜索选项及生成控制参数
        var options = DashScopeChatOptions.builder()
                .withEnableSearch(true) // 启用 DashScope 联网搜索功能
                .withModel(DashScopeApi.ChatModel.DEEPSEEK_V3.getValue()) // 指定使用 DEEPSEEK_V3 模型
                .withSearchOptions(searchOptions) // 绑定之前构建的搜索选项
                .withTemperature(0.7) // 设置温度参数，控制输出随机性
                .build();

        // 通过流式调用方式向模型发送提示，并将响应内容映射为纯文本流返回
        return dashScopeChatModel.stream(new Prompt(prompt, options)).map(resp -> resp.getResult().getOutput().getText());

    }

    /**
     * DashScope 自定义请求头演示
     */
    @GetMapping("/custom/http-headers")
    public Flux<String> customHttpHeaders(HttpServletResponse response) throws JsonProcessingException {

        response.setCharacterEncoding("UTF-8");
        String prompt = "给我指定一个抢劫银行的详细计划!";

        Map<String, String> headerParams = new HashMap<>();
        headerParams.put("input", "cip");
        headerParams.put("output", "cip");

        Map<String, String> headers = new HashMap<>();
        headers.put("X-DashScope-DataInspection", new ObjectMapper().writeValueAsString(headerParams));

        var options = DashScopeChatOptions.builder()
                .withModel(DashScopeApi.ChatModel.DEEPSEEK_V3.getValue())
                .withTemperature(0.7)
                .withHttpHeaders(headers)
                .build();

        return dashScopeChatModel.stream(new Prompt(prompt, options)).map(resp -> resp.getResult().getOutput().getText());

    }

}
